# data_loading.py
import json
import pandas as pd
import numpy as np
import torch

# 设置随机种子
def set_seed(seed):
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

set_seed(42)

# 加载数据
def load_data(file_path):
    """读取JSON格式数据"""
    with open(file_path, 'r', encoding='utf-8') as f:
        data = [json.loads(line) for line in f.readlines()]
    return pd.DataFrame(data)

# 加载证型数据及映射
def load_syndrome_mapping(syndrome_file):
    """创建证型名称到ID的映射"""
    syndrome_df = load_data(syndrome_file)
    syndrome_to_id = {s['Name']: i for i, s in enumerate(syndrome_df.to_dict('records'))}
    id_to_syndrome = {v: k for k, v in syndrome_to_id.items()}
    return syndrome_df, syndrome_to_id, id_to_syndrome

# 加载固定训练集和验证集（新增完整函数）
def load_fixed_datasets(train_file, dev_file, syndrome_file):
    """加载已划分好的训练集和验证集"""
    print(f"加载训练集: {train_file}")
    train_df = load_data(train_file)
    print(f"加载验证集: {dev_file}")
    dev_df = load_data(dev_file)
    print(f"加载证型数据: {syndrome_file}")
    syndrome_df, syndrome_to_id, id_to_syndrome = load_syndrome_mapping(syndrome_file)

    # 确保训练集和验证集中的证型都在映射中
    train_syndromes = set(train_df['syndrome'].unique())
    dev_syndromes = set(dev_df['syndrome'].unique())
    all_syndromes = set(syndrome_to_id.keys())

    missing_train = train_syndromes - all_syndromes
    missing_dev = dev_syndromes - all_syndromes

    if missing_train:
        print(f"警告: 训练集中有{len(missing_train)}个证型未在syndrome.json中定义")
    if missing_dev:
        print(f"警告: 验证集中有{len(missing_dev)}个证型未在syndrome.json中定义")

    return train_df, dev_df, syndrome_df, syndrome_to_id, id_to_syndrome

# 示例使用（仅测试）
if __name__ == "__main__":
    train_df, dev_df, _, _, _ = load_fixed_datasets('train.json',
                                                    'dev.json',
                                                    'syndrome.json')
    print(f"训练集: {len(train_df)}条, 验证集: {len(dev_df)}条")